import tushare as ts
from DB.comm.StockDaily import DbHandler
from DB.config import Base
from DB.comm.model import StockDaily
from DB.comm.BaseDb import BaseDb
from datetime import timedelta
import time


ts.set_token('5e9928d0932a3e94db90443dff9dc19e771d6c2ef2a5fbc1e10577a7')

#设置你的token
# df = pro.user(token='5e9928d0932a3e94db90443dff9dc19e771d6c2ef2a5fbc1e10577a7')

pro = ts.pro_api()
db = BaseDb("root", "123456", "127.0.0.1", "3306", "test")
stock_daily = DbHandler(db)

if __name__ == "__main__":
    if 0:
        ts_codes = stock_daily.get_all_ts_codes()
        i = 0
        for ts_code in ts_codes:
            last_row = stock_daily.get_last_stock_daily(ts_code)
            t =  (last_row.trade_date + timedelta(days=1)).strftime("%Y%m%d")
            print( ts_code, t)

            while True:
                try:
                    df = pro.daily(ts_code=ts_code, start_date=t)

                    df = df.sort_values("trade_date")  # 默认升序
                    print(df.head())
                    # exit()
                    i += 1
                    for index, row in df.iterrows():
                        # print(index, row['ts_code'], row["trade_date"], row["open"], row["high"], row["low"], row["close"], row["vol"], row["amount"])
                        stock_daily.add_stock_daily(StockDaily(ts_code=row['ts_code'], trade_date=row["trade_date"], open=row["open"], high=row["high"], low=row["low"], close=row["close"], pre_close=row["pre_close"], change=row["change"], pct_chg=row["pct_chg"], vol=row["vol"], amount=row["amount"]))
                    # # exit()
                    if i % 50 == 0:
                        print(f"已插入 {i} 条股票数据,休眠 60 秒后继续...")
                        time.sleep(60)
                    # print(df)
                    break   # 成功获取就跳出循环
                except Exception as e:
                    print(f"获取 {ts_code} 数据时出错: {e}，休眠 60 秒后重试...")
                    time.sleep(60)
    else:
        import datetime
        ts_codes = stock_daily.get_all_ts_codes()
        # 使用集合而不是列表，提高查找效率（O(1) vs O(n)）
        ts_codes_set = set(ts_codes)


        today = datetime.date.today()
        # today = datetime.date(2025, 9, 26)  # 指定日期为2025年1月1日
        today_str = today.strftime("%Y%m%d")
        latest_trade_date = datetime.date(2025, 10, 27)
        latest_trade_date_str = latest_trade_date.strftime("%Y%m%d")
        print(f"拉取数据从 {latest_trade_date_str} 到 {today_str}")
        # 遍历从 latest_trade_date 到 today 的日期
        current_date = latest_trade_date
        total_inserted = 0
        
        while current_date <= today:
            date_str = current_date.strftime("%Y%m%d")
            print(f"拉取 {date_str} 的数据")
            try:
                # 获取该日期的日线数据
                daily_df = pro.daily(trade_date=date_str)
                print(f"daily_df 列数: {len(daily_df.columns)}")
                
                if not daily_df.empty:
                    date_inserted = 0
                    for index, row in daily_df.iterrows():
                        # 使用集合查找，效率更高
                        if row['ts_code'] in ts_codes_set:
                            # 使用 ignore_duplicate 方法，避免重复数据导致的异常
                            stock_daily.add_stock_daily_ignore_duplicate(
                                StockDaily(
                                    ts_code=row['ts_code'], 
                                    trade_date=row["trade_date"], 
                                    open=row["open"], 
                                    high=row["high"], 
                                    low=row["low"], 
                                    close=row["close"], 
                                    pre_close=row["pre_close"], 
                                    change=row["change"], 
                                    pct_chg=row["pct_chg"], 
                                    vol=row["vol"], 
                                    amount=row["amount"]
                                )
                            )
                            date_inserted += 1

                    print(f"已处理 {date_str} 的数据，新增 {date_inserted} 条记录")
                    total_inserted += date_inserted
                else:
                    print(f"{date_str} 无交易数据")
                
                # 移动到下一天
                current_date += datetime.timedelta(days=1)
            except Exception as e:
                print(f"获取 {date_str} 数据时出错: {e}，休眠 10 秒后继续...")
                time.sleep(10)
            
        
        print(f"\n总计新增 {total_inserted} 条记录")
        # df = pro.daily(start_date=latest_trade_date_str)
        # for index, row in df.iterrows():
        #     if row['ts_code'] in ts_codes:
        #         # print(f"{index} {row['ts_code']}")
        #         stock_daily.add_stock_daily(StockDaily(ts_code=row['ts_code'], trade_date=row["trade_date"], open=row["open"], high=row["high"], low=row["low"], close=row["close"], pre_close=row["pre_close"], change=row["change"], pct_chg=row["pct_chg"], vol=row["vol"], amount=row["amount"]))
            #     print(row['ts_code'], row["trade_date"], row["open"], row["high"], row["low"], row["close"], row["vol"], row["amount"])
            # print(row['ts_code'], row["trade_date"], row["open"], row["high"], row["low"], row["close"], row["vol"], row["amount"])
        # exit()
        # print(df)
        # print(f"df 条数: {len(df)}")
        # print(f"latest_trade_date: {stock_daily.get_latest_trade_date()}")

        # print(f"今天的日期: {today_str}")
        
        # print(f"ts_codes 条数: {len(ts_codes)}")
        # for ts_code in df['ts_code']:
        #     if ts_code not in ts_codes:
        #         print(ts_code)
        # print("================")
        # # print(df['ts_code'])
        # lt = list(df['ts_code'])
        # # print(lt)
        # for ts_code in ts_codes:
        #     if ts_code not in lt:
        #         print(ts_code)
        # print("================")

        # df = df.sort_values("trade_date")  # 默认升序
        # # print(df)
        # # exit()
        # i += 1
        # for index, row in df.iterrows():
        #     print(index, row['ts_code'], row["trade_date"], row["open"], row["high"], row["low"], row["close"], row["vol"], row["amount"])
        #     stock_daily.add_stock_daily(StockDaily(ts_code=row['ts_code'], trade_date=row["trade_date"], open=row["open"], high=row["high"], low=row["low"], close=row["close"], pre_close=row["pre_close"], change=row["change"], pct_chg=row["pct_chg"], vol=row["vol"], amount=row["amount"]))
        # # exit()
        # if i % 50 == 0:
        #     print(f"已插入 {i} 条股票数据")
        # print(df)

    # 获取指定股票从 start_date 开始的日线数据
    # df = pro.daily(ts_code='000001.SZ', start_date='20250829')

    # print(df)